"""
@Author: Fhz
@Create Date: 2024/1/18 15:26
@File: xgboost.py
@Description: 
@Modify Person Date: 
"""
import numpy as np
import xgboost as xgb
from sklearn.metrics import accuracy_score, f1_score, recall_score, precision_score
import datetime
import pickle
from time import time


if __name__ == '__main__':
    y_path = ["../trajectory_prediction/Resave_dataset/y_test_1.npy",
              "../trajectory_prediction/Resave_dataset/y_test_2.npy",
              "../trajectory_prediction/Resave_dataset/y_train_1.npy",
              "../trajectory_prediction/Resave_dataset/y_train_2.npy",
              "../trajectory_prediction/Resave_dataset/y_train_3.npy",
              "../trajectory_prediction/Resave_dataset/y_train_4.npy",
              "../trajectory_prediction/Resave_dataset/y_train_5.npy",
              "../trajectory_prediction/Resave_dataset/y_train_6.npy"]

    x_path = ["../trajectory_prediction/Resave_dataset/X_test1.npy",
              "../trajectory_prediction/Resave_dataset/X_test2.npy",
              "../trajectory_prediction/Resave_dataset/X_train1.npy",
              "../trajectory_prediction/Resave_dataset/X_train2.npy",
              "../trajectory_prediction/Resave_dataset/X_train3.npy",
              "../trajectory_prediction/Resave_dataset/X_train4.npy",
              "../trajectory_prediction/Resave_dataset/X_train5.npy",
              "../trajectory_prediction/Resave_dataset/X_train6.npy"]

    y_path1 = ["../trajectory_prediction/Resave_dataset/y_valid_1.npy",
               "../trajectory_prediction/Resave_dataset/y_valid_2.npy"]

    x_path1 = ["../trajectory_prediction/Resave_dataset/X_valid1.npy",
               "../trajectory_prediction/Resave_dataset/X_valid2.npy"]

    for i in range(len(x_path)):
        X_tmp = np.load(file=x_path[i])
        y_tmp = np.load(file=y_path[i])

        if i == 0:
            X_train = X_tmp
            y_train = y_tmp
        else:
            X_train = np.vstack([X_train, X_tmp])
            y_train = np.vstack([y_train, y_tmp])

    for i in range(len(x_path1)):

        X_tmp = np.load(file=x_path1[i])
        y_tmp = np.load(file=y_path1[i])

        if i == 0:
            X_test = X_tmp
            y_test = y_tmp
        else:
            X_test = np.vstack([X_test, X_tmp])
            y_test = np.vstack([y_test, y_tmp])

    X_train = X_train[:, 10:, :]
    X_test = X_test[:, 10:, :]

    X_train = X_train.reshape((len(X_train), -1))
    X_test = X_test.reshape((len(X_test), -1))

    dtrain = xgb.DMatrix(X_train, y_train)  # 特征矩阵和标签都进行一个传入
    dtest = xgb.DMatrix(X_test, y_test)

    time0 = time()

    param = {'verbosity': 1
            , 'objective': 'multi:softmax'
            , 'num_class': 3
            , "subsample": 1
            , "max_depth": 8
            , "eta": 0.15
            , "gamma": 1
            , "lambda": 1
            , "alpha": 0
            , "colsample_bytree": 1
            , "colsample_bylevel": 1
            , "colsample_bynode": 1
            , "tree_method": "gpu_hist"
            , "gpu_id": 1
            }

    num_round = 150  # n_estimators ,

    bst = xgb.train(param, dtrain, num_round)
    model_name = "lstm-xgboost.dat"
    print(model_name)
    pickle.dump(bst, open(model_name, "wb"))

    preds = bst.predict(dtest)
    print("accuracy_score:")
    print(accuracy_score(y_test, preds))

    print("==========================")

    print("precision_score 0:")
    print(precision_score(preds, y_test, labels=[0], average="macro"))
    print("precision_score 1:")
    print(precision_score(preds, y_test, labels=[1], average="macro"))
    print("precision_score 2:")
    print(precision_score(preds, y_test, labels=[2], average="macro"))

    print("==========================")

    print("recall_score 0:")
    print(recall_score(preds, y_test, labels=[0], average="macro"))
    print("recall_score 1:")
    print(recall_score(preds, y_test, labels=[1], average="macro"))
    print("recall_score 2:")
    print(recall_score(preds, y_test, labels=[2], average="macro"))

    print("==========================")

    print("f1_score 0:")
    print(f1_score(preds, y_test, labels=[0], average="macro"))
    print("f1_score 1:")
    print(f1_score(preds, y_test, labels=[1], average="macro"))
    print("f1_score 2:")
    print(f1_score(preds, y_test, labels=[2], average="macro"))

    print(datetime.datetime.fromtimestamp(time() - time0).strftime("%M:%S:%f"))

